import os, sys
import glob
from PIL import Image
import argparse

def txtLabel_to_xmlLabel(source_txt_path,source_img_path,save_xml_path):
    if not os.path.exists(save_xml_path):
        os.makedirs(save_xml_path)
    for file in os.listdir(source_txt_path):
        print(file)
        img_path = os.path.join(source_img_path,file.replace('.txt','.jpg')) #png to jpg
        img_file = Image.open(img_path)
        txt_file = open(os.path.join(source_txt_path,file)).read().splitlines()
        # print(txt_file)
        xml_file = open(os.path.join(save_xml_path,file.replace('.txt','.xml')), 'w')
        width, height = img_file.size
        xml_file.write('<annotation>\n')
        xml_file.write('\t<folder>' + 'UAV_data' + '</folder>\n')
        xml_file.write('\t<filename>' + img_path + '</filename>\n')
        xml_file.write('\t<source>\n')
        xml_file.write('\t\t<database>The UAV autolanding</database>\n')
        xml_file.write('\t\t<annotation>UAV AutoLanding</annotation>\n')
        xml_file.write('\t\t<image>flickr</image>\n')
        xml_file.write('\t\t<flickrid>NULL</flickrid>\n')
        xml_file.write('\t</source>\n')
        xml_file.write('\t<size>\n')
        xml_file.write('\t\t<width>'+ str(width) + '</width>\n')
        xml_file.write('\t\t<height>'+ str(height) + '</height>\n')
        xml_file.write('\t\t<depth>' + str(3) + '</depth>\n')
        xml_file.write('\t</size>\n')
        xml_file.write('\t\t<segmented>0</segmented>\n')

        for line in txt_file:
            print(line)
            line_split = line.split(' ')
            classes = str(line_split[0])
            xmin = float(line_split[4])
            ymin = float(line_split[5])
            xmax = float(line_split[6])
            ymax = float(line_split[7])

            xml_file.write('\t<object>\n')
            xml_file.write('\t\t<name>'+ classes +'</name>\n')
            xml_file.write('\t\t<pose>Unspecified</pose>\n')
            xml_file.write('\t\t<truncated>0</truncated>\n')
            xml_file.write('\t\t<difficult>0</difficult>\n')
            xml_file.write('\t\t<bndbox>\n')
            xml_file.write('\t\t\t<xmin>' + str(int(round(xmin))) + '</xmin>\n')
            xml_file.write('\t\t\t<ymin>' + str(int(round(ymin))) + '</ymin>\n')
            xml_file.write('\t\t\t<xmax>' + str(int(round(xmax))) + '</xmax>\n')
            xml_file.write('\t\t\t<ymax>' + str(int(round(ymax))) + '</ymax>\n')
            xml_file.write('\t\t</bndbox>\n')
            xml_file.write('\t</object>\n')
        xml_file.write('</annotation>')

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--source_txt_path', type=str, default="labels")
    parser.add_argument('--source_img_path', type=str, default="images")
    parser.add_argument('--save_xml_path', type=str, default="save_xml_path")
    opt = parser.parse_args()

    txtLabel_to_xmlLabel(opt.source_txt_path,opt.source_img_path,opt.save_xml_path)


# python rope2d_to_xml.py --source_txt_path label --source_img_path training_image --save_xml_path training_xml